package edu.kit.csl.pisa.models;

import java.util.Arrays;

import edu.kit.csl.pisa.io.Logger;
import edu.kit.csl.pisa.ui.Configuration;

/*
This file is part of the PISA Alignment Tool.

Copyright (C) 2013
Karlsruhe Institute of Technology
Cognitive Systems Lab (CSL)
Felix Stahlberg

PISA is free software: you can redistribute it and/or modify
it under the terms of the GNU General Public License as published by
the Free Software Foundation, either version 3 of the License, or
(at your option) any later version.

PISA is distributed in the hope that it will be useful,
but WITHOUT ANY WARRANTY; without even the implied warranty of
MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE. See the
GNU General Public License for more details.

You should have received a copy of the GNU General Public License
along with PISA. If not, see <http://www.gnu.org/licenses/>.
*/

/**
 * Negative binomial sentence length distributions. Separate r,p parameters for
 * each source sentence length.
 */
public class NBSentenceLengthModel implements SentenceLengthModel {

	protected Model4P alignmentModel;
	
	// Model parameters
	protected int[] r;
	protected double[] logP;
	protected double[] logPInv; // log (1-p)
		
	protected NBEstimator estimators[];
	
	final int maxSrcSenLen;

	/**
	 * Sole constructor. 
	 * 
	 * @param model alignment model
	 */
	public NBSentenceLengthModel(Model4P model) {
		this.alignmentModel = model;
		maxSrcSenLen = Configuration.getSingleton().getInteger("maxSrcSenLen");
		double logHalf = Math.log(0.5);
		r = new int[maxSrcSenLen + 1];
		logP = new double[maxSrcSenLen + 1];
		logPInv = new double[maxSrcSenLen + 1];
		estimators = new NBEstimator[maxSrcSenLen + 1];
		Arrays.fill(logP, logHalf);
		Arrays.fill(logPInv, logHalf);
		for (int i = 0; i <= maxSrcSenLen; i++) {
			r[i] = i; // Initialize to mean since p = 1-p
			estimators[i] = new MLENBEstimator(i);
		}
		r[0] = 1;
	}

	/**
	 * Get the probability according the difference between source
	 * and target sentence length.
	 * {@inheritDoc}
	 */
	@Override
	public double getProbability(int nSrcWords, int nTrgtWords) {
		int thisR = r[nSrcWords];
		double acc = nTrgtWords * logP[nSrcWords] 
				+ thisR * logPInv[nSrcWords]
				- AlignmentModel.logFacFert[nTrgtWords];
		int maxI = nTrgtWords + thisR - 1;
		for (int i = thisR; i <= maxI; i++) {
			acc += Math.log(i);
		}
		
		return acc;
	}

	/* (non-Javadoc)
	 * @see SentenceLengthModel#count(double, int, int)
	 */
	@Override
	public void count(double weight, int nSrcWords, int nTrgtWords) {
		estimators[nSrcWords].addTrainingCase(weight, nTrgtWords);
	}

	/**
	 * Use the alignment model given in the constructor to normalize the
	 * collected fractional counts.
	 * 
	 * @see #count(double, int, int)
	 */
	@Override
	public void normalize() {
		Logger log = Logger.getSingleton();
		for (int i = 0; i <= maxSrcSenLen; i++) {
			estimators[i].estimate();
			double curP = estimators[i].getP();
			int curR = estimators[i].getR();
			r[i] = curR;
			logP[i] = Math.log(curP);
			logPInv[i] = Math.log(1.0 - curP);
			estimators[i].clear();
			log.debug("NegBinom Length Model: srcLen=" + i + ", r=" +
				curR + ", p=" + curP + ", mu=" + (curP * curR / (1.0 - curP)));
		}
	}

	/**
	 * Uses {@link AlignmentModel#dumpParameter(String, Object[])} to dump
	 * the p and r parameters to the file system.
	 * 
	 * @param prefix prefix of the dump file to create
	 * @param postfix postfix of the dump file to create
	 * @see SentenceLengthModel#dumpToFilesystem(String, String)
	 */
	@Override
	public void dumpToFilesystem(String prefix, String postfix) {
		double[][] arr = new double[1][];
		arr[0] = logP;
		alignmentModel.dumpParameter(prefix + "p" + postfix, arr);
		int[][] arrInt = new int[1][];
		arrInt[0] = r;
		alignmentModel.dumpParameter(prefix + "r" + postfix, arr);
	}
}
